from fastapi import FastAPI, File, UploadFile
import io
import numpy as np
from PIL import Image
from io import BytesIO
from magic_pdf.model.re_table.get_table_cell import ModelRunner  # 假设 ModelRunner 在 model_runner.py 文件中
import logging
import uvicorn
app = FastAPI()

# 初始化模型加载器
model_runner = ModelRunner()

@app.post("/predict_label/")
async def predict_label(file: UploadFile = File(...)):
    """
    接收图像文件并返回表格标签（'wired_table' 或 'wireless_table'）。
    """
    image_data = await file.read()
    image = Image.open(BytesIO(image_data))
    image = np.array(image)

    # 获取表格标签
    label = model_runner.get_table_label(image)
    
    return {"label": label}

@app.post("/predict_cells/")
async def predict_cells(file: UploadFile = File(...), label_name: str = "auto", threshold: float = 0.3):
    """
    接收图像文件并返回表格单元格的坐标。
    :param label_name: 表格标签 ('wired_table' 或 'wireless_table' 或 'auto')
    :param threshold: 用于预测时过滤的置信度阈值，默认为0.3
    """
    image_data = await file.read()
    image = Image.open(BytesIO(image_data))
    image = np.array(image)

    # 获取单元格坐标
    cell_coordinates = model_runner.get_cell_coordinates(image, label_name, threshold)
    
    return {"cell_coordinates": cell_coordinates}
# 启动 FastAPI 应用
if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=8000)